from collections import deque
from math import sqrt

from khronos.utils import INF
from khronos.statistics.ftable import FTable
from khronos.statistics.plotter import get_plotter

class Tally(object):
    def __init__(self, numeric=True, storing=True):
        self.__numeric = bool(numeric)
        self.__storing = bool(storing)
        if self.__storing:
            self.__values = deque()
            self.__weights = deque()
        else:
            self.__values = None
            self.__weights = None
        self.__min = INF
        self.__max = -INF
        self.__sum = 0.0
        self.__sumsquares = 0.0
        self.__wsum = 0.0
        self.__wsumsquares = 0.0
        self.__ftable = FTable()
        self.__wftable = FTable()
        
    def clear(self, numeric=None, storing=None):
        if numeric is None: numeric = self.__numeric
        if storing is None: storing = self.__storing
        self.__init__(numeric=numeric, storing=storing)
        
    def collect(self, value, weight=1.0):
        """Register a value with a specified weight in the tally collector. If no weight is 
        specified, the default value of 1 is used."""
        if self.__storing:
            self.__values.append(value)
            self.__weights.append(weight)
        if self.__numeric:
            squared_value = value * value
            self.__min = min(self.__min, value)
            self.__max = max(self.__max, value)
            self.__sum += value
            self.__sumsquares += squared_value
            self.__wsum += value * weight
            self.__wsumsquares += squared_value * weight
        self.__ftable.collect(value, 1.0)
        self.__wftable.collect(value, weight)
        
    def _add_weight(self, value, weight):
        """Used by TSeries to add weight information only after closing an interval."""
        if self.__storing:
            assert value is self.__values[-1]
            self.__weights[-1] += weight
        self.__wsum += weight * value
        self.__wsumsquares += weight * value * value
        self.__wftable.collect(value, weight)
        
    def merge(self, tally):
        if not isinstance(tally, Tally):
            raise TypeError("Tally expected")
        if self.__storing and not tally.__storing:
            raise ValueError("storing tally cannot merge with non-storing tally")
        if self.__numeric != tally.__numeric:
            raise ValueError("numeric tally cannot merge with non-numeric tally")
        if self.__storing:
            self.__values.extend(tally.__values)
            self.__weights.extend(tally.__weights)
        if self.__numeric:
            self.__min = min(self.__min, tally.__min)
            self.__max = max(self.__max, tally.__max)
            self.__sum += tally.__sum
            self.__sumsquares += tally.__sumsquares
            self.__wsum += tally.__wsum
            self.__wsumsquares += tally.__wsumsquares
        self.__ftable.merge(tally.__ftable)
        self.__wftable.merge(tally.__wftable)
        
    def report(self, display=True):
        fmt = "%20s = %s"
        lines = [object.__repr__(self)]
        lines.append(fmt % ("count", self.count())) 
        lines.append(fmt % ("min", self.min())) 
        lines.append(fmt % ("max", self.max())) 
        lines.append(fmt % ("sum", self.sum())) 
        lines.append(fmt % ("wsum", self.wsum())) 
        lines.append("-" * 20)
        lines.append(fmt % ("mean", self.mean())) 
        lines.append(fmt % ("var", self.var()))
        lines.append(fmt % ("std_dev", self.std_dev())) 
        lines.append("-" * 20)
        lines.append(fmt % ("wmean", self.wmean()))
        lines.append(fmt % ("wvar", self.wvar()))
        lines.append(fmt % ("wstd_dev", self.wstd_dev()))
        report = "\n".join(lines)
        if display:
            print report
        else:
            return report
            
    # -----------------------------------------------
    # Basic properties
    @property
    def ftable(self):
        return self.__ftable
        
    @property
    def wftable(self):
        return self.__wftable
        
    @property
    def storing(self):
        return self.__storing
        
    @property
    def numeric(self):
        return self.__numeric
        
    # -----------------------------------------------
    # Container methods
    def __len__(self):
        return int(self.__ftable.total())
        
    def __getitem__(self, idx):
        if not self.__storing:
            raise ValueError("unsubscriptable object (non-storing)")
        return self.__values[idx], self.__weights[idx]
        
    def __iter__(self):
        if not self.__storing:
            raise ValueError("uniterable object (non-storing)")
        idx = 0
        while idx < len(self.__values):
            yield self.__values[idx], self.__weights[idx]
            idx += 1
            
    def iter_values(self):
        for value, _ in self:
            yield value
            
    def iter_weights(self):
        for _, weight in self:
            yield weight
            
    # -----------------------------------------------
    # Basic statistics
    def count(self):
        return self.__ftable.total()
        
    def min(self):
        return self.__min
        
    def max(self):
        return self.__max
        
    def sum(self):
        return self.__sum
        
    def mean(self):
        count = self.__ftable.total()
        if count == 0.0:
            return 0.0
        return self.__sum / count
        
    def var(self):
        count = self.__ftable.total()
        if count == 0.0:
            return 0.0
        return self.__sumsquares / count - (self.__sum / count) ** 2
        
    def std_dev(self):
        return sqrt(self.var())
        
    def wsum(self):
        return self.__wsum
        
    def wmean(self):
        weight = self.__wftable.total()
        if weight == 0.0:
            return 0.0
        return self.__wsum / weight
        
    def wvar(self):
        weight = self.__wftable.total()
        if weight == 0.0:
            return 0.0
        return self.__wsumsquares / weight - (self.__wsum / weight) ** 2
        
    def wstd_dev(self):
        return sqrt(self.wvar())
        
    # -----------------------------------------------
    # Plotting
    def box_plot(self, target=None, **kwargs):
        if self.__storing:
            values = self.__values
        else:
            values = []
            for value, occurrences in self.__ftable.iteritems():
                values.extend([value] * occurrences)
        return get_plotter(target).box_plot(values, **kwargs)
        